import json
from make_heatmap import make_all_heatmap,group_point,make_group_heatmap
import cv2
import os
import numpy as np
from tqdm import tqdm
def resize_point(image,point,size):
    ori_size = image.shape[:-1]
    zoom_x = size[0]/ori_size[0]
    zoom_y = size[1]/ori_size[1]
    image = cv2.resize(image,size[::-1])
    point = [[p[0]*zoom_x,p[1] * zoom_y] for p in point]
    return image,point

def save_heatmap(json_dir,data_path,heatmap_path,resize_image_path,size,sigma = None):
    os.makedirs(heatmap_path,exist_ok=True)
    os.makedirs(resize_image_path,exist_ok=True)
    json_data = json.load(open(json_dir))
    for index in tqdm(range(len(json_data))):
        image = cv2.imread(os.path.join(data_path,json_data[index]['name'] + '.jpg'))
        label = json_data[index]['point']
        label = [p[0][::-1] for p in label]
        image,label  = resize_point(image,label,size)

        for point in label:
            cv2.circle(image, (int(point[1]), int(point[0])), 2, (255, 128, 255), thickness=5)
        
        image_shape = size
        # label = 255*make_all_heatmap(image_shape[1],image_shape[0],anno=label,sigma=sigma)
        label = group_point(label)
        label = make_group_heatmap(image_shape[1], image_shape[0], anno=label, sigma=sigma)
        np.save(os.path.join(heatmap_path,json_data[index]['name'] + '.npy'),label.astype(np.uint8))
        # cv2.imwrite(os.path.join(resize_image_path,json_data[index]['name'] + '.jpg'),image)
        
        
        # # 将heatmap进行可视化操作，用来检查
        # cv2.imshow('image',image)
        # show_label = (np.sum(label,axis=0)).astype(np.uint8)
        # cv2.imshow('heatmap',show_label)
        # cv2.waitKey()
        #
        

def adjust_json(data_path,json_dir,new_json_dir):
    ## 对json数据进行处理，避免出现部分错位标注。
    json_data = json.load(open(json_dir))
    new_json_data = []
    for index in tqdm(range(len(json_data))):
        image = cv2.imread(os.path.join(data_path, json_data[index]['name'] + '.jpg'))
        label = json_data[index]['point']
        label = [p[0][::-1] for p in label]
        image_shape = image.shape[:-1]
        # print(label)
        error = 0
        for num,p in enumerate(label):
            if p[0] < 0 or p[0] > image_shape[0] or p[1] < 0 or p[1] > image_shape[1]:
                error += 1
                print('error location')
                json_data[index]['point'][num][0]= [0,0]
        new_json_data.append(json_data[index])
    with open(new_json_dir,'w') as f:
        json.dump(new_json_data,f)
    
if __name__ == '__main__':
    adjust_json(data_path='./data/image',
                json_dir='label_file/point.json',
                new_json_dir='label_file/new_point.json')
    # save_heatmap(json_dir='label_file/point.json',
    #              data_path='./data/image',
    #              heatmap_path='./data/heatmap_640_512_group_01',
    #              resize_image_path='./data/image_650_512',
    #              size=(640,512),
    #              sigma=[9,11,17,25])